package handler

import (
	"context"
	"errors"
	"gitee.com/simple-set/simple.io/src/process/codec"
	"gitee.com/simple-set/simple.io/src/process/socks"
	"gitee.com/simple-set/simple.io/src/socket"
	"io"
	"log/slog"
	"net"
	"time"
)

type SocksBindHandler struct{}

func NewSocksBindHandler() *SocksBindHandler {
	return &SocksBindHandler{}
}

func (b *SocksBindHandler) Bind(request *socks.Request, response *socks.Response) error {
	sock, err := b.listen(request, response)
	if err != nil {
		return err
	}
	defer func() {
		_ = sock.Close()
	}()
	go func() {
		_, _ = codec.NewReadByteBuf(sock).WriteTo(request.Buff())
	}()
	if _, err := request.Buff().WriteTo(sock.Writer); err != nil {
		return err
	}
	return io.EOF
}

// Bind指令触发新的TCP监听, 只允许指定的ip地址连接且等待10秒超时
func (b *SocksBindHandler) listen(request *socks.Request, response *socks.Response) (*socket.Socket, error) {
	tcpServer := socket.NewTcpServer("")
	if err := tcpServer.Bind(); err != nil {
		_ = response.Reply(socks.NetworkUnreachable)
		return nil, err
	}
	defer func() {
		_ = tcpServer.Close()
	}()

	listenAddr := tcpServer.Listener().Addr().(*net.TCPAddr)
	response.BndAddr = listenAddr.IP
	response.BndPort = uint16(listenAddr.Port)
	if err := response.Reply(socks.SuccessReply); err != nil {
		return nil, err
	}

	sockets := make(chan *socket.Socket)
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
	go func() {
		for i := 0; i < 5; i++ {
			if sock, err := tcpServer.Listen(); err != nil {
				slog.Error("Bind TCP listener exception", "err", err)
				break
			} else {
				sockets <- sock
			}
		}
		cancel()
	}()

	for {
		select {
		case sock := <-sockets:
			sockAddr := sock.RemoteAddr().(*net.TCPAddr)
			if sockAddr.IP.Equal(request.GetDstAddr().IP) {
				_ = response.Reply(socks.SuccessReply)
				return sock, nil
			} else {
				slog.Warn("Connection address mismatch", "expectedAddr", request.GetDstAddrString(), "requestAddr", sockAddr)
				_ = sock.Close()
			}
		case <-ctx.Done():
			_ = response.Reply(socks.TtlExpired)
			return nil, errors.New("BndAddr waiting for timeout")
		}
	}
}
